#!/usr/bin/python
# -*- coding: utf-8 -*-
# @version        : 1.0
# @Create Time    : 2022/11/23 11:21 
# @File           : initialize.py
# @IDE            : PyCharm
# @desc           : 简要说明
import json
from sqlalchemy import insert
from core.database import db_getter
from utils.excel.excel_manage import ExcelManage
from application.settings import BASE_DIR, VERSION
import os
from apps.dispatch.taskpool import models as task_pool_models
from apps.dispatch.assignment import models as task_assign_models
from apps.dispatch.feedback import models as task_feedback_models
from apps.dispatch.rule import models as task_rule_models
from apps.check.param import models as check_param_models
from apps.check.feature import models as check_feature_models
from apps.check.vendor import models as check_vendor_models
from apps.check.sample import models as check_sample_models

import subprocess


class InitializeData:
    """
    初始化数据

    生成步骤：
        1. 读取数据
        2. 获取数据库
        3. 创建数据
    """

    SCRIPT_DIR = os.path.join(BASE_DIR, 'scripts', 'initialize')

    def __init__(self):
        self.sheet_names = []
        self.datas = {}
        self.ex = None
        self.db = None
        self.__serializer_data()
        self.__get_sheet_data()

    @classmethod
    def migrate_model(cls):
        """
        模型迁移映射到数据库
        """
        subprocess.check_call(['alembic', 'stamp', 'head', '--purge'], cwd=BASE_DIR)
        subprocess.check_call(['alembic', 'revision', '--autogenerate', '-m', f'{VERSION}'], cwd=BASE_DIR)
        subprocess.check_call(['alembic', 'upgrade', 'head'], cwd=BASE_DIR)
        print(f"{VERSION} 数据库表迁移完成")

    def __serializer_data(self):
        """
        序列化数据，将excel数据转为python对象
        """
        self.ex = ExcelManage()
        self.ex.open_workbook(os.path.join(self.SCRIPT_DIR, 'data', 'init.xlsx'), read_only=True)
        self.sheet_names = self.ex.get_sheets()

    def __get_sheet_data(self):
        """
        获取工作区数据
        """
        for sheet in self.sheet_names:
            sheet_data = []
            self.ex.open_sheet(sheet)
            headers = self.ex.get_header()
            datas = self.ex.readlines(min_row=2, max_col=len(headers))
            for row in datas:
                sheet_data.append(dict(zip(headers, row)))
            self.datas[sheet] = sheet_data

    async def __generate_data(self, table_name: str, model):
        """
        生成数据

        :param table_name: 表名
        :param model: 数据表模型
        """
        async_session = db_getter()
        db = await async_session.__anext__()
        datas = []
        for data in self.datas.get(table_name):
            for key, value in data.items():
                if isinstance(value, str) and value.startswith("[") and value.endswith("]"):
                    try:
                        data[key] = json.loads(value)
                    except Exception:
                        raise
            datas.append(data)
        print("table_name: ", table_name)
        
        try:
            await db.execute(insert(model), datas)
            await db.flush()
            await db.commit()
        except Exception as e:
            print("insert error:", e)
            await db.rollback()
        print(f"{table_name} 表数据已生成")

    async def generate_dispatch(self):
        """
        生成调度相关生产任务测试数据
        """
        await self.__generate_data("dc_dispatch_rule", task_rule_models.DCRule)
        await self.__generate_data("dc_dispatch_task", task_pool_models.DCTask)
        await self.__generate_data("dc_dispatch_feedback", task_feedback_models.DCTaskFeedback)
        await self.__generate_data("dc_dispatch_assign_users", task_assign_models.DCTaskAssignUsers)
        await self.__generate_data("dc_check_vendor", check_vendor_models.DCVendor)
        await self.__generate_data("dc_check_param", check_param_models.DCParam)
        await self.__generate_data("dc_check_asr_feature", check_feature_models.DCASRFeature)
        await self.__generate_data("dc_check_sample_group", check_sample_models.DCSampleGroup)
        await self.__generate_data("dc_check_sample_standard", check_sample_models.DCSampleStandard)

    async def run(self):
        """
        执行初始化工作
        """
        self.migrate_model()
        await self.generate_dispatch()
        print(f"{VERSION} 数据已初始化完成")
